# =====================================================================================
#
#                   QLoRA SFT 微调执行脚本
#
#   本脚本假设训练数据已被预处理并保存为 'sft_prepared_dataset.json'
#   该文件应由 utils.py 或其他数据准备脚本生成。
#
# =====================================================================================

# -------------------------------------------------------------------------------------
# 步骤 1: 导入所有必要的库
# -------------------------------------------------------------------------------------
import torch
from datasets import load_dataset
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    TrainerCallback,
    EvalPrediction,
)
from peft import LoraConfig
from trl import SFTConfig, SFTTrainer
import os

# -------------------------------------------------------------------------------------
# 步骤 2: 定义模型、预处理好的数据集路径和其他全局变量
# -------------------------------------------------------------------------------------
# 模型 ID
model_name = "Qwen/Qwen2-1.5B-Instruct"
# !! 关键 !!: 指向已准备好的数据集文件
prepared_dataset_path = "sft_prepared_dataset.json"
# 微调后模型的保存名称
new_model_name = "qwen2-1.5b-sft-qlora-from-prepared-data"
# !! 重要 !!: 替换为你的 Hugging Face 用户名
hub_model_id = f"your-hf-username/{new_model_name}"

# -------------------------------------------------------------------------------------
# 步骤 3: 加载预处理好的数据集
# -------------------------------------------------------------------------------------
print(f"--- 步骤 3: 加载预处理好的数据集 ---")
if not os.path.exists(prepared_dataset_path):
    raise FileNotFoundError(
        f"错误: 未找到数据集文件 '{prepared_dataset_path}'。\n"
        "请先运行数据准备脚本 (例如使用 utils.py 中的函数) 来生成此文件。"
    )

# SFTTrainer 需要 'messages' 列，我们假设预处理脚本已生成此格式
dataset = load_dataset("json", data_files=prepared_dataset_path, split="train")

if "messages" not in dataset.column_names:
     raise ValueError("数据集必须包含一个名为 'messages' 的列。请检查数据准备脚本。")

# 划分训练集和评估集
split_dataset = dataset.train_test_split(test_size=0.1, seed=42)
train_dataset = split_dataset["train"]
eval_dataset = split_dataset["test"]
print(f"数据集加载完毕。训练样本: {len(train_dataset)}, 评估样本: {len(eval_dataset)}")


# -------------------------------------------------------------------------------------
# 步骤 4: 配置 QLoRA (4-bit 量化) 和 PEFT (LoRA)
# -------------------------------------------------------------------------------------
print(f"\n--- 步骤 4: 配置 QLoRA 和 PEFT ---")
quantization_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
    bnb_4bit_use_double_quant=True,
)

peft_config = LoraConfig(
    r=16,
    lora_alpha=32,
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
    target_modules="all-linear",
)

# -------------------------------------------------------------------------------------
# 步骤 5: 加载模型和分词器
# -------------------------------------------------------------------------------------
print(f"\n--- 步骤 5: 加载模型和分词器 ({model_name}) ---")
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    quantization_config=quantization_config,
    device_map="auto",
    trust_remote_code=True
)

# -------------------------------------------------------------------------------------
# 步骤 6: (可选) 定义自定义函数和回调
# -------------------------------------------------------------------------------------
def compute_metrics(p: EvalPrediction):
    return {"loss": p.metrics["eval_loss"]}

class LoggingCallback(TrainerCallback):
    def on_step_end(self, args, state, control, **kwargs):
        if state.is_local_process_zero and state.global_step % args.logging_steps == 0 and state.log_history:
            last_log = state.log_history[-1]
            loss = last_log.get('loss') or last_log.get('train_loss', 'N/A')
            print(f"Step: {state.global_step}, Loss: {loss}")

# -------------------------------------------------------------------------------------
# 步骤 7: 配置 SFTConfig (包含所有参数)
# -------------------------------------------------------------------------------------
print(f"\n--- 步骤 7: 配置 SFTConfig (训练参数) ---")
training_args = SFTConfig(
    output_dir=new_model_name,
    overwrite_output_dir=True,
    do_train=True,
    do_eval=True,
    eval_strategy="steps",
    eval_steps=100,
    per_device_train_batch_size=4,
    per_device_eval_batch_size=8,
    gradient_accumulation_steps=2,
    num_train_epochs=1.0,
    learning_rate=2e-4,
    optim="paged_adamw_8bit",
    lr_scheduler_type="cosine",
    warmup_ratio=0.03,
    weight_decay=0.001,
    max_grad_norm=0.3,
    logging_strategy="steps",
    logging_steps=25,
    report_to=["tensorboard"],
    push_to_hub=True,
    hub_model_id=hub_model_id,
    bf16=True,
    gradient_checkpointing=True,
    gradient_checkpointing_kwargs={'use_reentrant': False},
    max_seq_length=1024,
    packing=True,
    assistant_only_loss=True,
    seed=42,
    remove_unused_columns=True,
)

# -------------------------------------------------------------------------------------
# 步骤 8: 初始化 SFTTrainer
# -------------------------------------------------------------------------------------
print(f"\n--- 步骤 8: 初始化 SFTTrainer ---")
trainer = SFTTrainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    tokenizer=tokenizer,
    peft_config=peft_config,
    compute_metrics=compute_metrics,
    callbacks=[LoggingCallback()],
)

# -------------------------------------------------------------------------------------
# 步骤 9: 开始训练并保存
# -------------------------------------------------------------------------------------
print(f"\n--- 步骤 9: 开始 SFT 微调 ---")
trainer.train()

print("\n--- 训练完成！ ---")
final_model_path = f"{new_model_name}/final_model"
trainer.save_model(final_model_path)
print(f"最终模型已保存至: {final_model_path}")
print(f"模型将被/已推送到 Hugging Face Hub: https://huggingface.co/{hub_model_id}")